import cv2
import numpy as np
import os
from multiprocessing.dummy import Pool
from numba import jit


class DeHaze:
    def __init__(self) -> None:
        self.__img_path_list = []
        self.output_dir = ""
        self.pool = Pool()

    def read_files(self, img_path_list):
        """
        输入照片路径别表，甄别并填充路径列表到类私有属性中。
        @params img_path_list: string[]  航空影像路径列表
        @return: void
        """
        null_list = []
        for img_path in img_path_list:
            if os.path.exists(img_path):
                null_list.append(img_path)
        self.__img_path_list = null_list
        return self

    def read_dirs(self, dir_path):
        """
        输入照片目录，甄别并填充路径列表到类私有属性中。
        @params dir_path: string  航空影像总目录
        @return: void
        """
        null_list = []
        for root, dirs, files in os.walk(dir_path):
            for file in files:
                if (
                    os.path.splitext(file)[-1] == ".jpg"
                    or os.path.splitext(file)[-1] == ".JPG"
                ):
                    null_list.append(os.path.join(root, file).replace("\\", "/"))
        self.__img_path_list = null_list
        return self
    
    @jit(forceobj=True)
    def process(self, output_dir):
        """
        多进程处理影像去雾
        @params output_dir:string 输出目录
        @return：void
        """
        self.output_dir = output_dir
        self.pool.map(self.__singleImg, self.__img_path_list)

    @jit(forceobj=True)
    def __singleImg(self, img_path):
        m = self.__deHaze(cv2.imread(img_path) / 255.0) * 255
        m = self.__brightness(m, 2, 5)
        out_path = os.path.join(
            self.output_dir, img_path.split("/")[-2], img_path.split("/")[-1]
        ).replace("\\", "/")
        cv2.imwrite(out_path, m)
        print(out_path)

    # 私有方法
    @jit(forceobj=True)
    def __deHaze(self, m, r=81, eps=0.001, w=0.95, maxV1=0.80, bGamma=False):
        """
        去雾主算法程序
        @params m: <class 'numpy.ndarray'> 待处理的图像
        @params r: int 滤波半径
        @params eps：int 偏离值
        @params w
        @params maxV1： float
        @params bGamma: boolen (默认False) 是否进行gamma校正
        @return
        """
        Y = np.zeros(m.shape)
        V1, A = self.__getV1(m, r, eps, w, maxV1)
        for k in range(3):
            Y[:, :, k] = (m[:, :, k] - V1) / (1 - V1 / A)
        Y = np.clip(Y, 0, 1)
        if bGamma:
            Y = Y ** (np.log(0.5) / np.log(Y.mean()))
        return Y

    @jit(forceobj=True)
    def __brightness(self, img, contrast, brightness):
        """
        调整图像对比度和亮度
        @parmas img：<class 'numpy.ndarray'> 待处理的图像
        @parmas contrast：float 对比度
        @params brightness： float 亮度
        @params reslut：<class 'numpy.ndarray'> 处理完的图像
        """
        rows, cols, channels = img.shape
        blank = np.zeros([rows, cols, channels], img.dtype)
        reslut = cv2.addWeighted(img, contrast, blank, 1 - contrast, brightness)
        return reslut

    @jit(forceobj=True)
    def __getV1(self, m, r, eps, w, maxV1):
        """
        计算大气遮罩图像V1和光照值A
        @params m: <class 'numpy.ndarray'> 取值范围[0,1]的图像
        @params r: int 引导滤波的卷积核尺寸
        @params eps: int 偏离值
        @params w: int 限制系数
        @params maxV1  float 透明度
        @return V1,A 大气遮罩 和 光照值
        """
        V1 = np.min(m, 2)
        zmf = self.__zmMinFilterGray(V1, 7)
        V1 = self.__guidedfilter(V1, zmf, r, eps)
        bins = 2000
        ht = np.histogram(V1, bins)
        d = np.cumsum(ht[0]) / float(V1.size)
        for lmax in range(bins - 1, 0, -1):
            if d[lmax] <= 0.999:
                break
        A = np.mean(m, 2)[V1 >= ht[1][lmax]].max()
        V1 = np.minimum(V1 * w, maxV1)
        return V1, A

    @jit(forceobj=True)
    def __guidedfilter(self, I, p, r, eps):
        """
        引导滤波
        @params I：<class 'numpy.ndarray'> 大气遮罩图像
        @params p: <class 'numpy.ndarray'> 最小值滤波图像
        @parmas r: int 卷积核尺寸
        @params eps: float 偏离值
        @return <class 'numpy.ndarray'>
        """
        m_I = cv2.boxFilter(I, -1, (r, r))
        m_p = cv2.boxFilter(p, -1, (r, r))
        m_Ip = cv2.boxFilter(I * p, -1, (r, r))
        cov_Ip = m_Ip - m_I * m_p

        m_II = cv2.boxFilter(I * I, -1, (r, r))
        var_I = m_II - m_I * m_I

        a = cov_Ip / (var_I + eps)
        b = m_p - a * m_I

        m_a = cv2.boxFilter(a, -1, (r, r))
        m_b = cv2.boxFilter(b, -1, (r, r))
        return m_a * I + m_b

    @jit(forceobj=True)
    def __zmMinFilterGray(self, src, r=7):
        """''
        最小值滤波(腐蚀)，获取暗通道图像
        @params src: <class 'numpy.ndarray'> 图像array
        @params r：int 滤波半径
        @result： <class numpy.ndarray>
        """
        return cv2.erode(src, np.ones((2 * r - 1, 2 * r - 1)))

if __name__ == "__main__":
    dhz = DeHaze()
    dhz.read_dirs("D:/arcgisDatasource/test").process("D:/arcgisDatasource/out")
    
